# ------------------------------------------------------------------------------
# Copyright (c) HQU
# Licensed under the HQU License.
# Written by Wang Youjije (youjieWang@stu.hqu.edu.cn)
# Modified by Wang Youjije (youjieWang@stu.hqu.edu.cn)
# ------------------------------------------------------------------------------


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import torch
import torch.nn as nn


# 计算关节点的损失，通过heatmap loss
class JointsMSELoss(nn.Module):
    def __init__(self, use_target_weight):
        '''
        :param use_target_weight: 是一个布尔类型的值
        '''
        super(JointsMSELoss, self).__init__()
        self.criterion = nn.MSELoss(reduction='mean')
        self.use_target_weight = use_target_weight

    def forward(self, output, target, target_weight):
        # TODO 这边的target_weight不知道怎么回事
        batch_size = output.size(0)
        num_joints = output.size(1)  # 这个是热图的通道数和关节点的个数是一样的
        # torch.split(input, split_size_or_sections, dim), dim=1表示从第一个维度开始拆分，split_size_or_sections表示分为多少个
        # 这里就是把给个热图给拆分开来
        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
        loss = 0

        for idx in range(num_joints):
            # 这里的做法就是每一个heatmap展开成一行，然后分别去计算损失
            heatmap_pred = heatmaps_pred[idx].squeeze()
            heatmap_gt = heatmaps_gt[idx].squeeze()
            if self.use_target_weight:
                loss += 0.5 * self.criterion(
                    heatmap_pred.mul(target_weight[:, idx]),
                    heatmap_gt.mul(target_weight[:, idx])
                )
            else:
                loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)

        return loss / num_joints


# TODO 4/25 这边还要定义一个坐标损失
class JointsCoordMSELoss(nn.Module):
    def __init__(self):
        super(JointsCoordMSELoss, self).__init__()

    def forward(self):
        pass
